from pydantic.v1 import BaseModel, Field
from crewai_tools import BaseTool
from typing import Type, List
import requests

class ArticleInfo(BaseModel):
    title: str
    link: str

class ArticleContent(BaseModel):
    article: str

class CSDNSearchToolInput(BaseModel):
    word: str = Field(..., title="Search Word", description="The word to search on CSDN")


class CSDNSearchTool(BaseTool):
    name: str = "CSDN Search Tool"
    description: str = "Search Article on CSDN"
    args_schema: Type[BaseModel] = CSDNSearchToolInput

    def _run(self, word: str) -> List[ArticleInfo]:
        search_list = []
        num = 1
        page = 1
        url = 'https://so.csdn.net/api/v3/search'
        data = {
            'q': word,
            'p': page,
        }
        headers = {
            'user-agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.0.0 Safari/537.36'
        }
        response = requests.get(url=url, params=data, headers=headers)
        for index in response.json()['result_vos']:
            if num == 5:
                break
            title = index["title"].replace('<em>', '').replace('</em>', '')
            author = index["nickname"].replace('<em>', '').replace('</em>', '')
            item = ArticleInfo(title=title, link=index['url'])
            num += 1
            search_list.append(item)
        return search_list

class CSDNViewArticleToolInput(BaseModel):
    link: str = Field(..., title="Article Link", description="The link of the article to view")

class CSDNViewArticleTool(BaseTool):
    name: str = "CSDN View Article Tool"
    description: str = "View Article on CSDN"
    args_schema: Type[BaseModel] = CSDNViewArticleToolInput

    def _run(self, link: str) -> ArticleContent:
        headers = {
            'user-agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.0.0 Safari/537.36'
        }
        response = requests.get(url=link, headers=headers)
        # 解析文章内容，blog-content-box类的div标签为文章内容
        html = response.text
        from bs4 import BeautifulSoup
        soup = BeautifulSoup(html, 'lxml')
        article = soup.find('div', class_='blog-content-box')
        return ArticleContent(article=article.text)

